from pypatnlp import *
from test import *
import numpy as np
from sklearn.decomposition import PCA
import matplotlib.pyplot as plt
import matplotlib
import sys
import os

def get_ai_paths():
    return [name for name in os.listdir(AI_PATH) if name.endswith('_ai.obj')]

#outf = sys.stdout
outf = open(os.path.join(TASK_PATH, 'output_005.csv'), 'wb')

outf.write('name,n,label,prediction\n')

for name in get_ai_paths():
    sys.stderr.write('Loading model\n')
    ai = None
    try:
        ai = load_obj(os.path.join(AI_PATH, name))
    except MemoryError, e:
        print e
        continue
    except EOFError, e:
        print e
        continue
    N, M = ai._corr_matrix.shape[0], ai._predict_corr_matrix.shape[0]
    print ai._corr_matrix.shape, ai._predict_corr_matrix.shape
    K = N + M
    X = np.vstack([ai._corr_matrix, ai._predict_corr_matrix])
    
    # do PCA
    pca = PCA(n_components=2)
    X = pca.fit_transform(X)
    svm = ai._svm
    
    ys = svm.predict(ai._predict_corr_matrix)
    
    name = name[:-7]
    outf.write('{0},{1},{2},{3}\n'.format(name, N, 'N', ys[0]))
    outf.write('{0},{1},{2},{3}\n'.format(name, N, 'Y', ys[1]))
    outf.flush()
    '''
    try:
        # figure out the contour of the `true zone` of oneclasssvm
        sys.stderr.write('Figuring out the contour\n')
        x_min, x_max = X[:, 0].min() - 25, X[:, 0].max() + 25
        y_min, y_max = X[:, 1].min() - 25, X[:, 1].max() + 25
        h = (x_max - x_min) / 100
        xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                            np.arange(y_min, y_max, h))
        ZX = np.vstack([xx.ravel(), yy.ravel()])
        ZX = np.transpose(ZX)
        Z = [0]
        # do the classification of grid row by row to avoid large memory
        # allocations due to inverse pca and large number of features
        for rowidx in range(ZX.shape[0]):
            row = pca.inverse_transform(ZX[rowidx:(rowidx+1),:])
            Z = np.append(Z, svm.predict(row))
        Z = Z[1:]
        Z = Z.reshape(xx.shape)
        
        # create the figure
        fig, ax = plt.subplots()
        
        plt.contourf(xx, yy, Z, colors=['0.75', '0.99'],)
        cont = plt.contour(xx, yy, Z, colors='0.1', linewidths=2, linestyles='solid')
        # draw the points
        trains = ax.scatter(X[:N,0], X[:N,1], marker='o', s=100, color='0.5') # training examples
        false = ax.scatter(X[N:(K-1),0], X[N:(K-1),1], marker='x', s=150, color='0.0') # false author
        true = ax.scatter(X[(K-1):,0], X[(K-1):,1], marker='o', s=150, color='0.0') # true author
        
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_xticks([])
        ax.set_yticks([])
        
        plt.legend([cont.collections[1], trains, true, false],
            ["learned frontier", "known articles",
            "article from the same author", "article from a different author"],
            loc="upper left",
            numpoints=1,
            prop=matplotlib.font_manager.FontProperties(size=11))

        fig.savefig(os.path.join(TASK_PATH, 'plots', name + '.pdf'))
        fig.savefig(os.path.join(TASK_PATH, 'plots', name + '.png'))
    except ValueError, e:
        print e
    '''

